package com.gitee.melin.bee.core.kubernetes;

import io.fabric8.kubernetes.client.OAuthTokenProvider;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.util.Base64;
import org.apache.commons.lang3.StringUtils;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.signer.Aws4Signer;
import software.amazon.awssdk.auth.signer.params.Aws4PresignerParams;
import software.amazon.awssdk.http.SdkHttpFullRequest;
import software.amazon.awssdk.http.SdkHttpMethod;
import software.amazon.awssdk.regions.Region;

// ref
// https://stackoverflow.com/questions/59853623/get-authentication-token-from-aws-eks-using-the-aws-java-sdk-v2/63264256#63264256
public class AwsEksOAuthTokenProvider implements OAuthTokenProvider {

    private static final String EKS_AUTH_HEADER = "x-k8s-aws-id";

    private static final String EKS_AUTH_PREFIX = "k8s-aws-v1.";

    private static final String STS_SERVICE_NAME = "sts";

    private static final String STS_HOST_PREFIX = "sts.";

    private static final String STS_PATH = "/";

    private static final String STS_ACTION = "GetCallerIdentity";

    private static final String STS_VERSION = "2011-06-15";

    private final String region;

    private final String accessKey;

    private final String secretKey;

    private final String k8sClusterName;

    public AwsEksOAuthTokenProvider(String region, String accessKey, String secretKey, String k8sClusterName) {
        this.region = region;
        this.accessKey = accessKey;
        this.secretKey = secretKey;
        this.k8sClusterName = k8sClusterName;
    }

    @Override
    public String getToken() {
        try {
            // https://docs.aws.amazon.com/zh_cn/IAM/latest/UserGuide/id_credentials_temp_region-endpoints.html
            URI uri = null;
            if (StringUtils.startsWith(region, "cn-")) { // 中国区域
                uri = URI.create("https://" + STS_HOST_PREFIX + region + ".amazonaws.com.cn");
            } else {
                uri = URI.create("https://" + STS_HOST_PREFIX + region + ".amazonaws.com");
            }

            SdkHttpFullRequest requestToSign = SdkHttpFullRequest.builder()
                    .method(SdkHttpMethod.GET)
                    .uri(uri)
                    .appendHeader(EKS_AUTH_HEADER, k8sClusterName)
                    .appendRawQueryParameter("Action", STS_ACTION)
                    .appendRawQueryParameter("Version", STS_VERSION)
                    .encodedPath(STS_PATH)
                    .build();

            ZoneId zoneId = ZoneId.of("UTC");
            ZonedDateTime now = ZonedDateTime.ofInstant(Instant.now(), zoneId);
            // token 有效期最多只能15分钟，超过15分钟，token 认证失败
            ZonedDateTime expirationDate =
                    ZonedDateTime.ofInstant(now.plusMinutes(15).toInstant(), zoneId);
            Aws4PresignerParams presignerParams = Aws4PresignerParams.builder()
                    .awsCredentials(AwsBasicCredentials.create(accessKey, secretKey))
                    .signingRegion(Region.of(region))
                    .signingName(STS_SERVICE_NAME)
                    .expirationTime(expirationDate.toInstant())
                    .build();
            SdkHttpFullRequest signedRequest = Aws4Signer.create().presign(requestToSign, presignerParams);
            String encodedUrl = Base64.getUrlEncoder()
                    .withoutPadding()
                    .encodeToString(signedRequest.getUri().toString().getBytes(StandardCharsets.UTF_8));
            return (EKS_AUTH_PREFIX + encodedUrl);
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage(), e);
        }
    }
}
